#!/usr/bin/env python

import os
import time
import socket
import urllib2
import cookielib
import mechanize

from colors import * #@UnusedWildImport

class Shelob():

    def __init__(self):
        self.cookieJar = cookielib.LWPCookieJar()
        self.browser = mechanize.Browser()
        self.__setBrowserAttributes__()
        self.timeout = 15
        self.pageDepth = 9
        self.fileExtension = None
        self.timeStamp = False
        self.queries = []
        self.links = []
        self.pages = []
        self.engine = None
           
    def __setBrowserAttributes__(self):
        self.browser.set_handle_equiv(True)
        self.browser.set_handle_referer(True)
        self.browser.set_handle_robots(False)
        self.browser.set_handle_refresh(mechanize._http.HTTPRefreshProcessor(), max_time = 1)
        self.browser.addheaders = [('User-agent', 'Mozilla/5.0 (Windows; U; Windows NT 6.1; en-US; rv:1.9.2.16) Gecko/20110319 BTRS28621 Firefox/3.6.16')]
        self.browser.set_cookiejar(self.cookieJar)
        if not os.path.exists(os.getcwd() + '/downloads/'):
            os.mkdir(os.getcwd() + '/downloads')
        self.saveDirectory = os.getcwd() + '/downloads/'
    
    def addQuery(self, query):
        self.queries.append(query)
    
    def start(self):
        for query in self.queries:
            print bold+C+"[*] Executing query:", query, W
            self.__executeQuery__(query)
            self.__getLinks__()
            if len(self.links) > 0:
                self.__getFiles__(self.links)
            if len(self.pages) > 0:
                self.__getAllPages__()
            self.pages = []
        self.queries = []
        return len(os.listdir(self.saveDirectory))
    
    def setEngine(self, name):
        if name.lower() == 'google':
            self.engine = {'url': 'http://google.com', 'exclude': 'google', 'name': 'Google', 'form': 'q', 'nr': 0, 'page': '/search?q=', 'pageKey': 'start='}
            return self.engine['name']
        elif name.lower() == 'bing':
            self.engine = {'url': 'http://bing.com', 'exclude': 'bing', 'name': 'Bing', 'form': 'q', 'nr': 0, 'page': '/search?q=', 'pageKey': 'first='}
            return self.engine['name']
        else:
            return None
    
    def __getTimeStamp__(self):
        if self.timeStamp:
            return '_' + time.strftime('%H-%M-%S')
        else:
            return ''
    
    def __executeQuery__(self, query):
        self.browser.open(url = self.engine['url'], timeout = self.timeout)
        self.browser.select_form(nr = self.engine['nr'])
        self.browser.form[self.engine['form']] = query
        self.browser.submit()
        self.browser.response().read()
    
    def __getAllPages__(self):
        pageNumber = 2
        for page in self.pages:
            print bold+G+'[*] Page %d of %d, please wait...' % (pageNumber, len(self.pages)+1), W
            self.__nextPage__(page)
            self.__getLinks__()
            self.__getFiles__(self.links)
            pageNumber += 1
    
    def __nextPage__(self, page):
        self.browser.open(url = self.engine['url']+page, timeout = self.timeout)
        self.browser.response().read()
    
    def __getFiles__(self, links):
        for link in links:
            if link.endswith(self.fileExtension):
                print bold+'[*] Get file:', link, W
                fileName = link.split('/')[-1] + self.__getTimeStamp__()
                self.__getFile__(link, self.saveDirectory + fileName)
    
    def __getFile__(self, url, saveTo):
        ''' Download a remote file from url and writes it to "saveTo" '''
        try:
            headers = {'User-Agent' : 'Mozilla/5.0 (compatible; MSIE 9.0; Windows NT 6.0; Trident/5.0'}
            req = urllib2.Request(url, None, headers)
            remoteSite = urllib2.urlopen(req, timeout = self.timeout)
            file = open(saveTo, 'w')
            while True:
                buffer = remoteSite.read(2048)
                if not buffer:
                    break
                file.write(buffer)
            file.close()
        except urllib2.URLError:
            if os.path.exists(saveTo) and os.path.isfile(saveTo):
                os.remove(saveTo)
            print bold+R+'[!] Error: The file does not appear to actually exist'+W
        except socket.timeout:
            os.remove(saveTo)
            print bold+R+'[!] Warning: A download timed out, deleted partial file'+W
        except:
            if os.path.exists(saveTo) and os.path.isfile(saveTo):
                os.remove(saveTo)
            print bold+R+"[!] OMFG: The download failed in a particularly gruesome way"+W
    
    def __getLinks__(self):
        self.links = []
        for link in self.browser.links():
            if link.url.startswith('http') and link.url.find(self.engine['exclude']) == -1:
                self.links.append(link.url)
            elif link.url.startswith(self.engine['page']) and link.url.find(self.engine['pageKey']) != -1:
                if len(self.pages) < self.pageDepth and not link in self.pages:
                    self.pages.append(link.url)
                